深度学习系统13. 硬件加速的实现

本节课将利用前两节课学到的知识，快速过一下底层的具体实现，以为第三次作业打基础。C++ 和 CUDA 的底层实现文件在 src 目录下，而 Python 的 NDArray 后端高层实现在 needle/backend_ndarray 下。本次作业的目的就是创造 NDArray 这个加速库，替换原有的 numpy 实现。根目录下的 CMakeLists.txt 和 Makefile 是为了构建库所需要的，有兴趣也可以查阅，但非硬性要求

关键数据结构

观察 python/needle/backend_ndarray/ndarray.py 中 NDArray 类的代码。这个类实际上是一个一维数组，根据底层后端的不同，可以通过不同方法开辟一个存储中的连续区域，然后再根据其它参数将其转换成一个逻辑上的多维数组。NDArray 包括五个成员变量

shape，元组类型
strides，元组类型
offset，整型
device，BackendDevice 类型
handle，独立于 C++ 实现，对应于一个指向一维数组的指针

对 handle 的作用可能大家会觉得比较费解，那我们来跟踪一些操作的执行过程。考虑如下代码

from needle import backend_ndarray as nd

x = nd.NDArray([1, 2, 3], device=nd.cuda())
y = x + 1

初始化

对第一条语句，首先我们希望在 GPU 上创建一个一维数组，包含三个元素。对应的，可以看一下 NDArray 的 __init__ 构造函数。首先，程序会把输入的 python 列表转化为一个 numpy 的 array，然后走 elif isinstance(other, np.ndarray) 的分支。在这个分支里，先会调用 self.make(other.shape, device=device 这条语句，创建数据结构，然后调用具体 device 的 from_numpy 方法，将这个 numpy 的 ndarray 压缩到连续空间，再用压缩后的数组初始化这个 array 的 _handle 成员变量（实际上是一个内存拷贝）

对于 make 这个类方法，其签名为

@staticmethod
def make(shape, strides=None, device=None, handle=None, offset=0)

如果 handle 为 None，则会根据 shape 创建一个大小合适的空数组

if handle is None:
    array._handle = array.device.Array(prod(shape))

具体的 Array 实现由各个 device 确定，而各个 device 有自己对应的后端。在 numpy 后端中，只是返回一个空 ndarray

self.array = np.empty(size, dtype=np.float32)

对于 CPU 后端，具体实现在 src/ndarray_backend_cpu.cc 中，通过 pybind 暴露。

py::class_<AlignedArray>(m, "Array")
    .definit<size_t>(), py::return_value_policy::take_ownership
    .def("ptr", &AlignedArray::ptr_as_int)
    .def_readonly("size", &AlignedArray::size);

实际上调用的是 AlignedArray 的构造函数，它会对分配一个指针并开辟指定大小的内存空间

struct AlignedArray {
  AlignedArray(const size_t size) {
    int ret = posix_mmalign((void**) &ptr, ALIGNMENT, size * ELEM_SIZE);
    if (ret != 0) throw std::bad_alloc();
    this->size = size;
  }
  ~AlignedArray() { free(ptr); }
  size_t ptr_as_int() { return (size_t)ptr; }
  scalar_t* ptr;
  size_t size;
}

另一种构造数组的方法是 from_numpy()（也是上面示例代码中构造的方法）。在 C++ 后端中，实现为一个数组拷贝，从 numpy 数组拷贝到对应指针指向的区域（即 _handle 指向的区域）

对于 GPU 后端，具体实现在 src/ndarray_backend_cuda.cu 中，Array 对应于一个 CudaArray 结构体。构造函数使用 cudaMalloc 根据指定元素数量分配显存，from_numpy 的实现和 C++ 后端的实现大同小异

数组与标量相加

对于 y = x + 1，实际上是重载了 NDArray 的加法运算符，定义在 __add__ 方法中。该方法进一步调用 ewise_or_scalar，根据操作数类型（是否为 NDArray 类对象）来判断实际调用逐元素相加函数，还是与标量相加函数。这里是执行 ewise_func(self.compact()._handle, other, out._handle)。其中第一个参数是自身的数组，第二个参数 other 是标量（1），第三个参数是输出的结果数组。对于 GPU 后端，最后会调用如下函数

void ScalarAdd(const CudaArray& a, scalar_t val, CudaArray *out) {
  // CudaOneDim计算grid和block的launch size
  CudaDimm dim = CudaOneDim(out->size);
  ScalarAddKernel<<<dim.grid, dim.block>>>(a.ptr, val, out->ptr, out->size);
}

__global__ void ScalarAddKernel(const scalar_t *a, scalar_t val, 
                                scalar_t *out, size_t size) {
  size_t gid = blockIdx.x * blockDim.x + threadIdx.x;
  if (gid < size) out[gid] = a[gid] + val;
}

对于 C++ 后端，实现会更简洁一些

void ScalarAdd(const AlignedArray& a, scalar_t val, AlignedArray *out) {
  for (size_t i = 0; i < a.size; i++) {
    out->ptr[i] = a.ptr[i] + val;
  }
}

在上面两个实现中，计算时都显式分配了输入和输出数组

另一点要注意的是，如果需要操作 handle 数组，往往要先做一次 compact 操作，例如 self.compact()._handle。具体原因在下面解释

NDArray 对象也可以执行 numpy() 操作，将底层数组赋予一个 numpy 端侧对象。该操作也由 GPU/CPU 后端实现，核心想法是先把该对象的 strides 信息转换到 numpy 的 strides 信息，然后根据给定的布局创建 numpy 数组（GPU 后端要先做一个将数组拷贝到 host 的操作）

建议在做作业之前先仔细阅读两个后端实现的 C++/CUDA 代码

Strides 操作

这里我们需要再回顾一个问题：我们为什么需要 strides。其原因是，通过控制 strides 和 shape，同样的底层数组会展现出不一样的布局。假设对下面的一维数组

import numpy as np
x = nd.array([0, 1, 2, 3, 4, 5], device=nd.cpu_numpy())

我们想将它 reshape 成一个 2×3 的形状，其实只需要做如下操作

z = nd.NDArray.make(shape=(2, 3), strides=(3, 1), device=x.device, 
                    handle=x._handle, offset=0)

这里只是做了一个浅拷贝，两者底层指向的是同一块内存空间。通过 strides，我们有 z[i][j] = data[i * strides[0] + j * strides[1]] ，结合指定的 shape，就达到了 reshape 的目的。此时有

z = NDArray([[0, 1, 2],
             [3, 4, 5]])

如果我们想获得 z 的一个切片，比如 [[1, 2], [4, 5]]（记为 b），又该如何做？这个切片的大小为 2×2，所以新的 NDArray 对象有 shape = (2, 2)。strides 在这里不动，我们只需要把 offset 设为 1 即可达到我们的目的。这里有 b[i][j] = data[offset + i * strides[0] + j * strides[1]]

strides 也可以用来完成其他操作。例如，对转置，只需要把 strides 的两个元素对调即可（当然也要对应地去修改 shape）。对广播，只需要在要广播的位置在 strides 上加 0。例如如果要把 2×3 的张量广播到 2×3×4，在变换 shape 以后只需要把 strides 置为 (3, 1, 0)。（这里笔者有一个译文：这符合 numpy 的广播规则吗？）

需要注意的是，NDArray 的底层实现都是直接在一维数组上操作，这里默认所有元素都是分配到了一个连续区域。然而，在经过 strides 和 offset 操作以后，底层的数据不再是连续存储的。例如，如果要对前面 b 做一个加 1 的操作，如果直接在原数组上执行，那么所有 6 个元素都会被修改，而不是 b 切片得到的 4 个元素。因此，需要定义一个 compact 方法：先看 strides 和 shape 是否正常，如果不正常，将对应的数组元素复制到一个连续的区域。在 needle 中，对逐元素操作，我们总是会对 NDArray 做一个 compact 操作，但是也可以通过一些其它技术手段在非连续区域做计算，避免对 compact 的频繁调用

CUDA 加速

最后我们讨论如何加入 CUDA 算子。（这里一部分演示是根据 EWiseMul 和 ScalarMul 实现 EWiseDiv 和 ScalarDiv，基本就是照猫画虎，就不赘述了。注意修改完以后需要重新 make，以及最好创建单独的测试文件，来刷新 python session）

这部分作业中，needle/autograd.py 已经从原来的 numpy 换成了抽象的 array_api，张量也不再是 numpy 的多维数组，而是自己实现的 NDArray 类型对象